package dev.marsh.sniper.ar.filter;

import android.content.Context;
import android.util.Log;

import org.opencv.android.Utils;
import org.opencv.core.DMatch;
import org.opencv.core.KeyPoint;
import org.opencv.core.Mat;
import org.opencv.core.MatOfDMatch;
import org.opencv.core.MatOfKeyPoint;
import org.opencv.core.Point;
import org.opencv.features2d.DescriptorExtractor;
import org.opencv.features2d.DescriptorMatcher;
import org.opencv.features2d.FeatureDetector;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import dev.marsh.sniper.core.BuildConfig;

public final class ImageDetectionFilter implements Filter {
    private static final String TAG = "ImageDetectionFilter";

    // Features of the reference image
    private final MatOfKeyPoint mReferenceKeyPoints = new MatOfKeyPoint();
    // Descriptors of the reference image's features
    private final Mat mReferenceDescriptors = new Mat();

    // Features of the scene (the current frame)
    private final MatOfKeyPoint mSceneKeyPoints = new MatOfKeyPoint();
    // Descriptors of the scene's features
    private final Mat mSceneDescriptors = new Mat();

    // A grayscale version of the scene
    private final Mat mGraySrc = new Mat();
    // Tentative matches of scene features and reference features
    private final MatOfDMatch mMatches = new MatOfDMatch();

    // A feature detector, which finds features in images
    private final FeatureDetector mFeatureDetector =
            FeatureDetector.create(FeatureDetector.ORB);

    // A descriptor extractor, which creates descriptors of features
    private final DescriptorExtractor mDescriptorExtractor =
            DescriptorExtractor.create(DescriptorExtractor.ORB);

    // A descriptor matcher, which matches features based on their descriptors
    private final DescriptorMatcher mDescriptorMatcher =
            DescriptorMatcher.create(DescriptorMatcher.BRUTEFORCE_HAMMINGLUT);

    public ImageDetectionFilter(final Context context, final int referenceImageResourceID) throws IOException {
        // Load the reference image from the app's resources
        Mat referenceImageGray = Utils.loadResource(context,
                referenceImageResourceID,
                Imgcodecs.CV_LOAD_IMAGE_GRAYSCALE);

        // Detect the reference features and compute their descriptors
        mFeatureDetector.detect(referenceImageGray, mReferenceKeyPoints);
        mDescriptorExtractor.compute(referenceImageGray, mReferenceKeyPoints, mReferenceDescriptors);

        long total = mReferenceKeyPoints.total();
        long desTotal = mReferenceDescriptors.total();

        if (BuildConfig.DEBUG) {
            Log.e(TAG, "Sample keypoints: " + total);
            Log.e(TAG, "Sample descriptors: " + desTotal);
        }
    }

    @Override
    public boolean apply(final Mat src) {
        // Convert the scene to grayscale
        Imgproc.cvtColor(src, mGraySrc, Imgproc.COLOR_RGBA2GRAY);
        // Detect the scene features, compute their descriptors, and match the scene descriptors to reference descriptors
        mFeatureDetector.detect(mGraySrc, mSceneKeyPoints);
        mDescriptorExtractor.compute(mGraySrc, mSceneKeyPoints, mSceneDescriptors);
        mDescriptorMatcher.match(mSceneDescriptors, mReferenceDescriptors, mMatches);

        long total = mSceneKeyPoints.total();
        long desTotal = mSceneDescriptors.total();

        if (BuildConfig.DEBUG) {
            Log.e(TAG, "Dst keypoints: " + total);
            Log.e(TAG, "Dst descriptors: " + desTotal);
        }

        // Attempt to find the target image in the scene
        return findSceneCorners();
    }

    private boolean findSceneCorners() {
        final List<DMatch> matchesList = mMatches.toList();
        if (matchesList.size() < 4) {
            // There are too few matches to find
            return false;
        }

        final List<KeyPoint> referenceKeyPointsList = mReferenceKeyPoints.toList();

        // Calculate the max and min distances between keyPoints
        double maxDist = 0.0;
        double minDist = 100;
        for (final DMatch match : matchesList) {
            final double dist = match.distance;
            if (dist < minDist) {
                minDist = dist;
            }
            if (dist > maxDist) {
                maxDist = dist;
            }
        }

        final ArrayList<Point> goodMatches = new ArrayList<Point>();
        final double maxGoodMatchDist = 0.31 * maxDist;
        for (final DMatch match : matchesList) {
            if (match.distance < maxGoodMatchDist) {
                goodMatches.add(referenceKeyPointsList.get(match.trainIdx).pt);
            }
        }

        int size = goodMatches.size();

        if (BuildConfig.DEBUG) {
            Log.e(TAG, "GoodMatches: " + size);
        }
        return size >= 5;
    }
}
